#!/usr/bin/env python

# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
测试使用：调试摄像头位置使用！！
显示实时摄像头画面，终端输出当前摄像头索引，按 'n' 切换到下一个摄像头，按 'q' 退出。
在显示完所有摄像头后自动退出。
"""

import argparse
import logging
import cv2
from typing import Any, Dict, List

from lerobot.cameras.configs import ColorMode
from lerobot.cameras.opencv.camera_opencv import OpenCVCamera
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
from lerobot.cameras.realsense.camera_realsense import RealSenseCamera
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig

logger = logging.getLogger(__name__)

def find_all_opencv_cameras() -> List[Dict[str, Any]]:
    """Finds all available OpenCV cameras plugged into the system."""
    all_opencv_cameras_info: List[Dict[str, Any]] = []
    logger.info("Searching for OpenCV cameras...")
    try:
        opencv_cameras = OpenCVCamera.find_cameras()
        for cam_info in opencv_cameras:
            all_opencv_cameras_info.append(cam_info)
        logger.info(f"Found {len(opencv_cameras)} OpenCV cameras.")
    except Exception as e:
        logger.error(f"Error finding OpenCV cameras: {e}")
    return all_opencv_cameras_info

def find_all_realsense_cameras() -> List[Dict[str, Any]]:
    """Finds all available RealSense cameras plugged into the system."""
    all_realsense_cameras_info: List[Dict[str, Any]] = []
    logger.info("Searching for RealSense cameras...")
    try:
        realsense_cameras = RealSenseCamera.find_cameras()
        for cam_info in realsense_cameras:
            all_realsense_cameras_info.append(cam_info)
        logger.info(f"Found {len(realsense_cameras)} RealSense cameras.")
    except ImportError:
        logger.warning("Skipping RealSense camera search: pyrealsense2 library not found or not importable.")
    except Exception as e:
        logger.error(f"Error finding RealSense cameras: {e}")
    return all_realsense_cameras_info

def find_and_print_cameras(camera_type_filter: str | None = None) -> List[Dict[str, Any]]:
    """Finds available cameras and prints their information."""
    all_cameras_info: List[Dict[str, Any]] = []
    if camera_type_filter:
        camera_type_filter = camera_type_filter.lower()
    if camera_type_filter is None or camera_type_filter == "opencv":
        all_cameras_info.extend(find_all_opencv_cameras())
    if camera_type_filter is None or camera_type_filter == "realsense":
        all_cameras_info.extend(find_all_realsense_cameras())

    if not all_cameras_info:
        if camera_type_filter:
            print(f"未找到任何 {camera_type_filter} 摄像头")
        else:
            print("未找到任何可用摄像头 (OpenCV 或 RealSense)")
    else:
        print("\n--- 检测到的摄像头 ---")
        for i, cam_info in enumerate(all_cameras_info):
            print(f"摄像头 #{i}:")
            for key, value in cam_info.items():
                if key == "default_stream_profile" and isinstance(value, dict):
                    print(f"  {key.replace('_', ' ').capitalize()}:")
                    for sub_key, sub_value in value.items():
                        print(f"    {sub_key.capitalize()}: {sub_value}")
                else:
                    print(f"  {key.replace('_', ' ').capitalize()}: {value}")
            print("-" * 20)
    return all_cameras_info

def create_camera_instance(cam_meta: Dict[str, Any]) -> Any | None:
    """Create and connect to a camera instance based on metadata."""
    cam_type = cam_meta.get("type")
    cam_id = cam_meta.get("id")
    instance = None
    logger.info(f"正在准备 {cam_type} 摄像头 ID {cam_id}...")
    try:
        if cam_type == "OpenCV":
            cv_config = OpenCVCameraConfig(
                index_or_path=cam_id,
                color_mode=ColorMode.RGB,
            )
            instance = OpenCVCamera(cv_config)
        elif cam_type == "RealSense":
            rs_config = RealSenseCameraConfig(
                serial_number_or_name=cam_id,
                color_mode=ColorMode.RGB,
            )
            instance = RealSenseCamera(rs_config)
        else:
            print(f"未知摄像头类型: {cam_type} (ID: {cam_id})，跳过")
            return None
        if instance:
            logger.info(f"正在连接 {cam_type} 摄像头: {cam_id}...")
            instance.connect(warmup=False)
            return instance
    except Exception as e:
        print(f"连接或配置 {cam_type} 摄像头 {cam_id} 失败: {e}")
        if instance and instance.is_connected:
            instance.disconnect()
        return None

def test_camera(cam_instance, cam_info: Dict[str, Any], index: int, total_cameras: int) -> bool:
    """测试并显示单个摄像头画面，返回是否继续测试下一个摄像头"""
    window_name = f"摄像头 #{index} ({cam_info.get('type')} ID: {cam_info.get('id')})"
    try:
        print(f"\n正在显示摄像头 #{index}:")
        print(f"  类型: {cam_info.get('type')}")
        print(f"  ID: {cam_info.get('id')}")
        print("  - 按 'n' 切换到下一个摄像头")
        print("  - 按 'q' 退出程序")

        while True:
            frame = cam_instance.read()
            if frame is not None:
                cv2.imshow(window_name, frame)
            else:
                print(f"摄像头 #{index} 读取帧失败")
                break

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                return False  # 退出程序
            elif key == ord('n'):
                return True  # 切换到下一个摄像头

        return False  # 默认退出

    except Exception as e:
        print(f"摄像头 #{index} 发生错误: {e}")
        return False
    finally:
        if cam_instance and cam_instance.is_connected:
            cam_instance.disconnect()
        cv2.destroyWindow(window_name)

def main():
    parser = argparse.ArgumentParser(
        description="显示实时摄像头画面，按 'n' 切换摄像头，按 'q' 退出，显示完所有摄像头后自动退出"
    )
    parser.add_argument(
        "camera_type",
        type=str,
        nargs="?",
        default=None,
        choices=["realsense", "opencv"],
        help="指定要显示的摄像头类型 ('realsense' 或 'opencv')，省略则显示所有摄像头"
    )
    args = parser.parse_args()

    # 初始提示
    print("欢迎使用摄像头测试程序！")
    print("操作说明：请将鼠标光标点击弹出图像然后按照提示操作")
    print("  - 按 'n' 切换到下一个摄像头")
    print("  - 按 'q' 随时退出程序")
    print("正在搜索可用摄像头...\n")

    # 查找所有可用摄像头
    cameras = find_and_print_cameras(camera_type_filter=args.camera_type)
    if not cameras:
        return

    print(f"共检测到 {len(cameras)} 个摄像头，开始测试...\n")

    # 依次测试每个摄像头
    for idx in range(len(cameras)):
        cam_instance = create_camera_instance(cameras[idx])
        if cam_instance is None:
            print(f"跳过摄像头 #{idx}，无法连接")
            continue
        continue_testing = test_camera(cam_instance, cameras[idx], idx, len(cameras))
        if not continue_testing:
            break  # 用户按 'q' 或发生错误，退出循环

    print("\n所有摄像头测试完成或已退出")
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()